##########
## Plot
##########

plot.gsynth <- function(x, # a gsynth object
                        type = "gap", # type of the plot
                        xlim = NULL, ylim = NULL, # axes limits
                        xlab = NULL, ylab = NULL, # axes labels
                        legendOff = FALSE,
                        raw = "band", # show raw data in "counterfactual" mode
                        # ("none","band","all")
                        main = NULL, # whether to show the title
                        nfactors = NULL, # whose loadings to be plotted
                        id = NULL, # individual plot
                        ...){
  
  
  ##-------------------------------##
  ## Checking Parameters
  ##-------------------------------##  
  
  outcome <- NULL
  ATT <- NULL
  CI.lower <- NULL
  CI.upper <- NULL
  co5 <- NULL
  co95 <- NULL
  tr5 <- NULL
  tr95 <- NULL
  group <- NULL
  L1 <- NULL
  
  if (class(x)!="gsynth") {
    stop("Not a \"gsynth\" object.")
  }
  if (!type %in% c("gap","counterfactual","factors","loadings","raw")) {
    stop("\"type\" option misspecified.")
  }
  if (is.null(xlim)==FALSE) {
    if (is.numeric(xlim)==FALSE) {
      stop("Some element in \"xlim\" is not numeric.")
    } else {
      if (length(xlim)!=2) {
        stop("xlim must be of length 2.")
      }
    }
  }
  if (is.null(ylim)==FALSE) {
    if (is.numeric(ylim)==FALSE) {
      stop("Some element in \"ylim\" is not numeric.")
    } else {
      if (length(ylim)!=2) {
        stop("ylim must be of length 2.")
      }
    }
  }
  if (is.null(xlab)==FALSE) {
    if (is.character(xlab) == FALSE) {
      stop("\"xlab\" is not a string.")
    } else {
      xlab <- xlab[1]
    }   
  }
  if (is.null(ylab)==FALSE) {
    if (is.character(ylab) == FALSE) {
      stop("\"ylab\" is not a string.")
    } else {
      ylab <- ylab[1]
    }   
  }
  if (is.logical(legendOff) == FALSE & is.numeric(legendOff)==FALSE) {
    stop("\"legendOff\" is not a logical flag.")
  }
  if (type == "counterfactual") {
    if (! raw %in% c("none","band","all")) {
      stop("\"raw\" option misspecifed.") 
    }
    if (is.null(id)==FALSE) {
      if (length(id)>1) {
        stop("More than 1 element in \"id\".") 
      }
    } 
  }
  if (is.null(main)==FALSE) {
    if (is.character(main) == FALSE) {
      stop("\"main\" is not a string.")
    } else {
      main <- main[1]
    }   
  }
  if (is.null(nfactors)==FALSE) {
    if (is.numeric(nfactors)==FALSE) {
      stop("\"nfactors\" is not a positive integer.")
    } else {
      nfactors <- nfactors[1]
      if (nfactors%%1!=0 | nfactors<=0) {
        stop("\"nfactors\" is not a positive integer.")
      }  
    } 
  } 
  
  ##-------------------------------##
  ## Plotting
  ##-------------------------------##  
  
  
  Y.tr <- x$Y.tr
  Y.co <- x$Y.co
  Y.ct <- x$Y.ct
  tb <- x$est.att
  Yb <- x$Y.bar[,1:2] ## treated average and counterfactual average
  tr <- x$tr
  pre <- x$pre
  TT <- x$T
  T0 <- x$T0
  p <- x$p
  m <- x$m
  Ntr <- x$Ntr
  Nco <- x$Nco
  N <- x$N 
  force <- x$force
  F.hat <- x$factor
  L.tr <- x$lambda.tr
  if (!is.null(L.tr)) {
    r <- dim(L.tr)[2]
  } else {
    r <- 0
  }
  if (is.null(id)==TRUE) {
    id <- x$id.tr
  }
  
  ## parameters
  line.width <- c(1.2,0.5)
  
  ## type of plots
  if (type == "raw"| type == "counterfactual" | type == "factors" |
      x$DID == TRUE | length(id) == 1) {
    time <- x$time
    if (length(id) == 1) {
      time.bf <- time[T0[which(id == x$id.tr)]]
    } else {
      time.bf <- time[unique(T0)]
    }
    
  } else if (type == "gap")  { ## variable treatment timing
    time <- c(1:TT) - min(T0)
    time.bf <- 0 ## before treatment
  }
  
  ## periods to show
  if (length(xlim) != 0) {
    show <- which(time>=xlim[1]& time<=xlim[2])
  } else {
    show <- 1:length(time)
  }
  nT <- length(show) 
  
  ## legend on/off
  if (legendOff == TRUE) {
    legend.pos <- "none"
  } else {
    legend.pos <- "bottom"
  }
  
  ############  START  ###############
  
  if (type == "raw") {
    ## axes labels
    if (is.null(xlab)==TRUE) {
      xlab <- x$index[2]
    } else if (xlab == "") {
      xlab <- NULL
    }
    if (is.null(ylab)==TRUE) {
      ylab <- x$Yname
    } else if (ylab == "") {
      ylab <- NULL
    }
    pst <- (1 - x$pre)
    for (i in 1:Ntr){
      pst[T0[i],i] <- 1 ## paint the period right before treatment
    }
    time.pst <- c(pst[show,] * time[show])
    time.pst <- time.pst[which(c(pst[show,])==1)]
    Y.tr.pst <- c(Y.tr[show,])[which(pst[show,]==1)]
    id.tr.pst <- matrix(rep(1:Ntr,each=TT),TT,Ntr,byrow=FALSE)[show,]
    id.tr.pst <- c(id.tr.pst)[which(pst[show,]==1)]
    
    data <- cbind.data.frame("time" = c(rep(time[show], N), time.pst),
                             "outcome" = c(c(Y.tr[show,]),
                                           c(Y.co[show,]),
                                           Y.tr.pst),
                             "type" = c(rep("tr",(Ntr*nT)),
                                        rep("co",(Nco*nT)),
                                        rep("tr.pst",length(Y.tr.pst))),
                             "id" = c(rep(1:N,each = nT), id.tr.pst*(-1)))
    
    ## theme
    p <- ggplot(data) + xlab(xlab) +  ylab(ylab) +
      theme(legend.position = legend.pos,
            plot.title = element_text(size=20,
                                      hjust = 0.5,
                                      face="bold",
                                      margin = margin(10, 0, 10, 0)))
    
    
    
    if (x$DID==TRUE) {
      p <- p + geom_vline(xintercept=time.bf,colour="white",size = 2) #+
        #annotate("rect", xmin= time.bf, xmax= Inf,
        #         ymin=-Inf, ymax=Inf, alpha = .3) 
    }
    
    ## main
    p <- p + geom_line(aes(time, outcome,
                           colour = type,
                           size = type,
                           linetype = type,
                           group = id))
    
    ## legend
    set.limits = c("tr","tr.pst","co")
    set.labels = c("Treated (Pre)",
                   "Treated (Post)",
                   "Controls")
    set.colors = c("#FC8D6280","red","#99999950")
    set.linetypes = c("solid","solid","solid")
    set.linewidth = c(0.5, 0.5, 0.5)
    
    p <- p + scale_colour_manual(limits = set.limits,
                                 labels = set.labels,
                                 values =set.colors) +
      scale_linetype_manual(limits = set.limits,
                            labels = set.labels,
                            values = set.linetypes) +
      scale_size_manual(limits = set.limits,
                        labels = set.labels,
                        values = set.linewidth) +
      guides(linetype = guide_legend(title=NULL, ncol=3),
             colour = guide_legend(title=NULL, ncol=3),
             size = guide_legend(title=NULL, ncol=3))
    
    ## title
    if (is.null(main) == TRUE) {
      p <- p + ggtitle("Raw Data")
    } else if (main!="") {
      p <- p + ggtitle(main)
    }
    
    ## ylim
    if (is.null(ylim) == FALSE) {
      p <- p + coord_cartesian(ylim = ylim)
    }
    p
    
  } else if (type == "gap") { 
    
    if (length(id) == 1 & !(id[1] %in% x$id.tr)) { ## error
      cat(paste(id,"not in the treatment group"))
    } else { ## no error
      
      ## axes labels
      if (is.null(xlab) == TRUE) {
        if (x$DID == TRUE) {
          xlab <- x$index[2]
        } else {
          xlab <- paste("Time relative to Treatment")
        }
      } else if (xlab == "") {
        xlab <- NULL
      }
      if (is.null(ylab) == TRUE) {
        ylab <- "Coefficient"
      } else if (ylab == "") {
        ylab <- NULL
      }
      
      ## title
      if (length(id) == 1) { ## id specified
        maintext <- paste(x$index[1],"=",id) 
      }  else {
        maintext <- "Estimated Average Treatment Effect"
      } 
      
      ## contruct data for plotting
      if (is.null(x$est.att)==TRUE) { 
        cat("Uncertainty estimates not available.\n")
        if (length(id) == 1) { ## id specified
          data <- cbind.data.frame(time, x$eff)[show,]
          colnames(data) <- c("time","ATT")
        } else {
          data <- cbind.data.frame(time, ATT = x$att)[show,] 
        } 
      } else {
        if (length(id) == 1) { ## id specified
          id <- which(x$id.tr == id)
          tb <- x$est.ind[,,id]
          time.bf <- time[T0[id]] 
          colnames(tb) <- c("ATT", "S.E.", "CI.lower", "CI.upper","p.value") 
        }
        data <- cbind.data.frame(time, tb)[show,]
      }
      
      ## plotting
      p <- ggplot(data) +
        geom_vline(xintercept = time.bf, colour="white",size = 2) +
        geom_hline(yintercept = 0, colour="white",size = 2) +
        ## annotate("rect", xmin= time.bf, xmax= Inf,
        ##          ymin=-Inf, ymax=Inf, alpha = .1,
        ##          fill = "yellow") +
        xlab(xlab) +  ylab(ylab) +
        theme(legend.position = legend.pos,
              plot.title = element_text(size=20,
                                        hjust = 0.5,
                                        face="bold",
                                        margin = margin(10, 0, 10, 0)))
      
      
      ## point estimates
      p <- p + geom_line(aes(time, ATT), size = 1.2)
      
      ## confidence intervals
      if (is.null(x$est.att)==FALSE) {
        p <- p + geom_ribbon(aes(x = time, ymin=CI.lower, ymax=CI.upper),alpha=0.2)
      }
      
      ## title
      if (is.null(main) == TRUE) {
        p <- p + ggtitle(maintext)
      } else if (main!=""){
        p <- p + ggtitle(main)
      }
      
      ## ylim
      if (is.null(ylim) == FALSE) {
        p <- p + coord_cartesian(ylim = ylim)
      }
      p
    }  ## end of "gap" (in case of no id error)
    
    
  } else if (type=="counterfactual") {
    
    ## axes labels
    if (is.null(xlab)==TRUE) {
      xlab <- x$index[2]
    } else if (xlab == "") {
      xlab <- NULL
    }
    if (is.null(ylab)==TRUE) {
      ylab <- x$Yname
    } else if (ylab == "") {
      ylab <- NULL
    } 
    
    if (length(id)==1 & !(id[1]%in%x$id.tr)) { ## error
      
      cat(paste(id,"not in the treatment group"))
      
    } else { ## one treated unit case
      
      
      if (length(id) == 1 | length(x$id.tr) == 1) { ## one treated unit
        
        if (is.null(id) == TRUE) {
          id <- x$id.tr
        }
        maintext <- paste("Treated and Counterfactual (",id,")",sep="") 
        tr.info <- Y.tr[,which(id==x$id.tr)]
        ct.info <- Y.ct[,which(id==x$id.tr)] 
        if (raw == "none") { 
          data <- cbind.data.frame("time" = rep(time[show],2),
                                   "outcome" = c(tr.info[show],
                                                 ct.info[show]),
                                   "type" = c(rep("tr",nT),
                                              rep("ct",nT)))
          ## theme
          p <- ggplot(data) + xlab(xlab) +  ylab(ylab) +
            geom_vline(xintercept=time.bf,colour="white",size = 2) +
            #annotate("rect", xmin= time.bf, xmax= Inf,
            #         ymin=-Inf, ymax=Inf, alpha = .3) +
            theme(legend.position = legend.pos,
                  plot.title = element_text(size=20,
                                            hjust = 0.5,
                                            face="bold",
                                            margin = margin(10, 0, 10, 0))) 
          ## main
          p <- p + geom_line(aes(time, outcome,
                                 colour = type,
                                 size = type,
                                 linetype = type)) 
          ## legend
          set.limits = c("tr","ct")
          set.labels = c("Treated", "Estimated Y(0)")
          set.colors = c("black","#888888")
          set.linetypes = c("solid","longdash")
          set.linewidth = rep(line.width[1],2)
          p <- p + scale_colour_manual(limits = set.limits,
                                       labels = set.labels,
                                       values =set.colors) +
            scale_linetype_manual(limits = set.limits,
                                  labels = set.labels,
                                  values = set.linetypes) +
            scale_size_manual(limits = set.limits,
                              labels = set.labels,
                              values = set.linewidth) +
            guides(linetype = guide_legend(title=NULL, ncol=2),
                   colour = guide_legend(title=NULL, ncol=2),
                   size = guide_legend(title=NULL, ncol=2)) 
          
        } else if  (raw == "band") {
          
          Y.co.90 <- t(apply(Y.co, 1, quantile, prob=c(0.05,0.95))) 
          data <- cbind.data.frame("time" = rep(time[show],2),
                                   "outcome" = c(tr.info[show],
                                                 ct.info[show]),
                                   "type" = c(rep("tr",nT),
                                              rep("ct",nT)))
          
          data.band <- cbind.data.frame(time, Y.co.90)[show,]
          colnames(data.band) <- c("time","co5","co95")
          
          
          ## theme 
          p <- ggplot(data) + xlab(xlab) +  ylab(ylab) +
            geom_vline(xintercept=time.bf,colour="white",size = 2) +
            #annotate("rect", xmin= time.bf, xmax= Inf,
            #         ymin=-Inf, ymax=Inf, alpha = .3) +
            theme(legend.position = legend.pos,
                  plot.title = element_text(size=20,
                                            hjust = 0.5,
                                            face="bold",
                                            margin = margin(10, 0, 10, 0)))
          
          ## main
          p <- p + geom_line(aes(time, outcome,
                                 colour = type,
                                 size = type,
                                 linetype = type))
          
          ## band
          p <- p + geom_ribbon(data = data.band,
                               aes(ymin = co5, ymax = co95, x=time),
                               alpha = 0.15, fill = "steelblue")
          
          set.limits = c("tr","co.band","ct")
          set.labels = c("Treated", "Controls (5-95% Quantiles)",
                         "Estimated Y(0)")
          set.colors = c("red","#4682B480","steelblue")
          set.linetypes = c("solid","solid","longdash")
          set.linewidth = c(line.width[1],4,line.width[1])
          
          p <- p + scale_colour_manual(limits = set.limits,
                                       labels = set.labels,
                                       values =set.colors) +
            scale_linetype_manual(limits = set.limits,
                                  labels = set.labels,
                                  values = set.linetypes) +
            scale_size_manual(limits = set.limits,
                              labels = set.labels,
                              values = set.linewidth) +
            guides(linetype = guide_legend(title=NULL, ncol=3),
                   colour = guide_legend(title=NULL, ncol=3),
                   size = guide_legend(title=NULL, ncol=3))  
          
        } else if (raw == "all") { ## plot all the raw data
          
          data <- cbind.data.frame("time" = rep(time[show],(2 + Nco)),
                                   "outcome" = c(tr.info[show],
                                                 ct.info[show],
                                                 c(Y.co[show,])),
                                   "type" = c(rep("tr",nT),
                                              rep("ct",nT),
                                              rep("raw.co",(Nco * nT))),
                                   "id" = c(rep("tr",nT),
                                            rep("ct",nT),
                                            rep(c(x$id.co), each = nT)))
          
          ## theme
          p <- ggplot(data) + xlab(xlab) +  ylab(ylab) +
            geom_vline(xintercept=time.bf,colour="white",size = 2) +
            #annotate("rect", xmin= time.bf, xmax= Inf,
            #         ymin=-Inf, ymax=Inf, alpha = .3) +
            theme(legend.position = legend.pos,
                  plot.title = element_text(size=20,
                                            hjust = 0.5,
                                            face="bold",
                                            margin = margin(10, 0, 10, 0)))
          
          ## main
          p <- p + geom_line(aes(time, outcome,
                                 colour = type,
                                 size = type,
                                 linetype = type,
                                 group = id))
          
          ## legend
          set.limits = c("tr","raw.co","ct")
          set.labels = c("Treated","Controls","Estimated Y(0)")
          set.colors = c("black","#99999950","#888888")
          set.linetypes = c("solid","solid","longdash")
          set.linewidth = c(line.width[1],line.width[2],line.width[1])
          
          p <- p + scale_colour_manual(limits = set.limits,
                                       labels = set.labels,
                                       values =set.colors) +
            scale_linetype_manual(limits = set.limits,
                                  labels = set.labels,
                                  values = set.linetypes) +
            scale_size_manual(limits = set.limits,
                              labels = set.labels,
                              values = set.linewidth) +
            guides(linetype = guide_legend(title=NULL, ncol=3),
                   colour = guide_legend(title=NULL, ncol=3),
                   size = guide_legend(title=NULL, ncol=3))
          
        } 
        
      } else { # begin multiple treated unit case
        maintext <- "Treated and Counterfactual Averages"
        if (raw == "none") {
          data <- cbind.data.frame("time" = rep(time[show],2),
                                   "outcome" = c(Yb[show,1],
                                                 Yb[show,2]),
                                   "type" = c(rep("tr",nT),
                                              rep("co",nT))) 
          ## theme 
          p <- ggplot(data) + xlab(xlab) +  ylab(ylab) +
            geom_vline(xintercept=time.bf,colour="white",size = 2) +
            #annotate("rect", xmin= time.bf, xmax= Inf,
            #         ymin=-Inf, ymax=Inf, alpha = .3) +
            theme(legend.position = legend.pos,
                  plot.title = element_text(size=20,
                                            hjust = 0.5,
                                            face="bold",
                                            margin = margin(10, 0, 10, 0)))
          ## main
          p <- p + geom_line(aes(time, outcome,
                                 colour = type,
                                 size = type,
                                 linetype = type))
          
          ## legend
          set.limits = c("tr","co")
          set.labels = c("Treated Avaerge",
                         "Estimated Y(0) Average")
          set.colors = c("red","steelblue")
          set.linetypes = c("solid","longdash")
          set.linewidth = rep(line.width[1],2)
          p <- p + scale_colour_manual(limits = set.limits,
                                       labels = set.labels,
                                       values =set.colors) +
            scale_linetype_manual(limits = set.limits,
                                  labels = set.labels,
                                  values = set.linetypes) +
            scale_size_manual(limits = set.limits,
                              labels = set.labels,
                              values = set.linewidth) +
            guides(linetype = guide_legend(title=NULL, ncol=2),
                   colour = guide_legend(title=NULL, ncol=2),
                   size = guide_legend(title=NULL, ncol=2)) 
          
        } else if  (raw == "band") {
          
          Y.tr.90 <- t(apply(Y.tr, 1, quantile, prob=c(0.05,0.95)))
          Y.co.90 <- t(apply(Y.co, 1, quantile, prob=c(0.05,0.95)))
          
          data <- cbind.data.frame("time" = rep(time[show],2),
                                   "outcome" = c(Yb[show,1],
                                                 Yb[show,2]),
                                   "type" = c(rep("tr",nT),
                                              rep("co",nT)))
          
          data.band <- cbind.data.frame(time, Y.tr.90, Y.co.90)[show,]
          colnames(data.band) <- c("time","tr5","tr95","co5","co95")
          
          ## theme 
          p <- ggplot(data) + xlab(xlab) +  ylab(ylab) +
            geom_vline(xintercept=time.bf,colour="white",size = 2) +
            #annotate("rect", xmin= time.bf, xmax= Inf,
            #         ymin=-Inf, ymax=Inf, alpha = .3) +
            theme(legend.position = legend.pos,
                  plot.title = element_text(size=20,
                                            hjust = 0.5,
                                            face="bold",
                                            margin = margin(10, 0, 10, 0)))
          ## main
          p <- p + geom_line(aes(time, outcome,
                                 colour = type,
                                 size = type,
                                 linetype = type))
          ## band
          p <- p + geom_ribbon(data = data.band,
                               aes(ymin = tr5, ymax = tr95, x=time),
                               alpha = 0.15, fill = "red") +
            geom_ribbon(data = data.band,
                        aes(ymin = co5, ymax = co95, x=time),
                        alpha = 0.15, fill = "steelblue")
          
          set.limits = c("tr","co","tr.band","co.band")
          set.labels = c("Treated Avaerge",
                         "Estimated Y(0) Average",
                         "Treated 5-95% Quantiles",
                         "Controls 5-95% Quantiles")
          set.colors = c("red","steelblue","#FF000030","#4682B480")
          set.linetypes = c("solid","longdash","solid","solid")
          set.linewidth = c(rep(line.width[1],2),4,4)
          
          p <- p + scale_colour_manual(limits = set.limits,
                                       labels = set.labels,
                                       values =set.colors) +
            scale_linetype_manual(limits = set.limits,
                                  labels = set.labels,
                                  values = set.linetypes) +
            scale_size_manual(limits = set.limits,
                              labels = set.labels,
                              values = set.linewidth) +
            guides(linetype = guide_legend(title=NULL, ncol=2),
                   colour = guide_legend(title=NULL, ncol=2),
                   size = guide_legend(title=NULL, ncol=2)) 
          
        } else if (raw == "all") { ## plot all the raw data
          
          data <- cbind.data.frame("time" = rep(time[show],(2 + N)),
                                   "outcome" = c(Yb[show,1],
                                                 Yb[show,2],
                                                 c(Y.tr[show,]),
                                                 c(Y.co[show,])),
                                   "type" = c(rep("tr",nT),
                                              rep("co",nT),
                                              rep("raw.tr",(Ntr * nT)),
                                              rep("raw.co",(Nco * nT))),
                                   "id" = c(rep("tr",nT),
                                            rep("co",nT),
                                            rep(c(x$id.tr,x$id.co),
                                                each = nT))) 
          ## theme
          p <- ggplot(data) + xlab(xlab) +  ylab(ylab) +
            geom_vline(xintercept=time.bf,colour="white",size = 2) +
            #annotate("rect", xmin= time.bf, xmax= Inf,
            #         ymin=-Inf, ymax=Inf, alpha = .3) +
            theme(legend.position = legend.pos,
                  plot.title = element_text(size=20,
                                            hjust = 0.5,
                                            face="bold",
                                            margin = margin(10, 0, 10, 0))) 
          ## main
          p <- p + geom_line(aes(time, outcome,
                                 colour = type,
                                 size = type,
                                 linetype = type,
                                 group = id))
          ## legend
          set.limits = c("tr","co","raw.tr","raw.co")
          set.labels = c("Treated Avaerge",
                         "Estimated Y(0) Average",
                         "Treated Raw Data",
                         "Controls Raw Data")
          set.colors = c("red","steelblue","#FC8D6280","#99999950")
          set.linetypes = c("solid","longdash","solid","solid")
          set.linewidth = rep(line.width,each=2)
          
          p <- p + scale_colour_manual(limits = set.limits,
                                       labels = set.labels,
                                       values =set.colors) +
            scale_linetype_manual(limits = set.limits,
                                  labels = set.labels,
                                  values = set.linetypes) +
            scale_size_manual(limits = set.limits,
                              labels = set.labels,
                              values = set.linewidth) +
            guides(linetype = guide_legend(title=NULL, ncol=2),
                   colour = guide_legend(title=NULL, ncol=2),
                   size = guide_legend(title=NULL, ncol=2)) 
        }
        
        
      } # end multiple treated unit case
      
      ## title
      if (is.null(main) == TRUE) {
        p <- p + ggtitle(maintext)
      } else if (main!="") {
        p <- p + ggtitle(main)
      }
      
      ## ylim
      if (is.null(ylim) == FALSE) {
        p <- p + coord_cartesian(ylim = ylim)
      }
    }
    p
    
  } else if (type=="factors") {
    
    if (x$r.cv==0) {
      cat("No factors included in the model.\n")
    } else {
      ## axes labels
      if (is.null(xlab)==TRUE) {
        xlab <- x$index[2]
      } else if (xlab == "") {
        xlab <- NULL
      }
      if (is.null(ylab)==TRUE) {
        ylab <- "Estimate"
      } else if (ylab == "") {
        ylab <- NULL
      }
      ## title
      if (is.null(main) == TRUE) {
        main <- "Latent Factors"
      } else if (main=="") {
        main <- NULL
      }
      ## prapre data
      L.co<-x$lambda.co
      norm<-sqrt(diag(t(L.co)%*%L.co)/(x$N-x$Ntr))
      data <- cbind.data.frame("time" = rep(time[show],r),
                               "factor" = c(F.hat[show,])*rep(norm,each=nT),
                               "group" = as.factor(c(rep(1:r,each=nT))))
      ## theme
      p <- ggplot(data) + xlab(xlab) +  ylab(ylab) + ggtitle(main) +
        geom_hline(yintercept=0,colour="white",size = 2) +
        theme(legend.position = legend.pos,
              plot.title = element_text(size=20,
                                        hjust = 0.5,
                                        face="bold",
                                        margin = margin(10, 0, 10, 0)))  
      ## main plot
      p <- p + geom_line(aes(time, factor,
                             colour = group,
                             group = group,
                             linetype = group #change made
                             ), size = 1)
      ## legend
      p <- p + guides(colour = guide_legend(title="Factor(s)"), linetype = guide_legend(title="Factor(s)") )
      
      ## ylim
      if (is.null(ylim) == FALSE) {
        p <- p + coord_cartesian(ylim = ylim)
      }
      p
    }
    
  } else if (type=="loadings") {
    
    
    if (x$r.cv==0) {
      cat("No factors are included in the model.\n") 
    } else {
      ## number of loadings to be plotted
      if (is.null(nfactors)==TRUE) {
        nfactors<-min(x$r.cv,4) 
      } else if (nfactors>x$r.cv) {
        cat("Too many factors specified. ")
        nfactors<-min(x$r.cv,4) 
      }
      if (nfactors == 1) {
        cat("Loadings for the first factor are shown...\n")
      } else if (nfactors < x$r.cv) {
        cat(paste("Loadings for the first",nfactors,"factors are shown...\n"))
      }
      
      
      ## title
      if (is.null(main) == TRUE) {
        main <- "Factor Loadings"
      } else if (main=="") {
        main <- NULL
      }
      
      ## prepare data
      L.hat <- rbind(x$lambda.tr, x$lambda.co)
      Lname <- Llabel <- c()
      for (i in 1:r) {
        Lname<-c(Lname,paste("L",i,sep=""))
        Llabel<-c(Llabel,paste("Factor",i))
      }
      colnames(L.hat) <- Lname
      rownames(L.hat) <- c()
      data <- cbind.data.frame(L.hat,
                               "id"=c(x$id.tr, x$id.co),
                               "group"=as.factor(c(rep("Treated",Ntr),
                                                   rep("Control",Nco))))
      
      if (nfactors == 1) {
        p <- ggplot(data, aes(x=group, y=L1, fill = group)) +
          geom_boxplot(alpha = 0.7) +
          coord_flip() + guides(fill=FALSE) +
          xlab("") + ylab("Factor Loading")
      } else {
        
        if (x$Ntr < 5) {
          my_dens <- function(data, mapping, ...) {
            ggplot(data = data, mapping = mapping) +
              geom_density(..., fill = "gray", alpha = 0.7, color = "gray50")
          }
          p <- GGally::ggpairs(data, mapping = aes(color = group),
                       columns = 1:nfactors,
                       columnLabels = Llabel[1:nfactors],
                       diag = list(continuous = my_dens),
                       title = main)
        } else {
          my_dens <- function(data, mapping, ...) {
            ggplot(data = data, mapping = mapping) +
              geom_density(..., alpha = 0.7, color = NA)
          }
          p <- GGally::ggpairs(data, mapping = aes(color = group, fill = group),
                       columns = 1:nfactors,
                       columnLabels = Llabel[1:nfactors],
                       diag = list(continuous = my_dens),
                       title = main) +
            theme(plot.title = element_text(hjust = 0.5))
        }
      } 
      p + theme_classic()
    }
    
  }## fig: loadings
  
  
  
}